from __future__ import annotations
from dataclasses import dataclass, field, fields
from typing import Annotated, Optional
from langchain_core.runnables import RunnableConfig, ensure_config


@dataclass(kw_only=True)
class Configuration:
    chat_model: Annotated[str, {"__template_metadata__": {"kind": "llm"}}] = field(
        default="cohere:command-a-03-2025",
        metadata={
            "description": "The chat model to use for generating responses. Defaults to cohere:command-a-03-2025."
            "Should be in the form: provider:model-name."
        },
    )

    embedding_model: Annotated[
        str, {"__template_metadata__": {"kind": "embedding"}}
    ] = field(
        default="cohere:embed-multilingual-v3.0",
        metadata={
            "description": "The embedding model to use for generating embeddings. Defaults to cohere:embed-multilingual-v3.0."
            "Should be in the form: provider:model-name."
        },
    )

    search_top_k: Annotated[int, {"__template_metadata__": {"kind": "search"}}] = field(
        default=5,
        metadata={
            "description": "The number of top results to return from the search. Defaults to 5."
        },
    )

    @classmethod
    def from_runnable_config(
        cls, config: Optional[RunnableConfig] = None
    ) -> Configuration:
        """Load configuration from the given config"""
        config = ensure_config(config)
        configurable = config.get("configurable") or {}
        _fields = {f.name for f in fields(cls) if f.init}
        return Configuration(**{k: v for k, v in configurable.items() if k in _fields})
